//! 虚拟机管理
//! 
//! Firecracker VM的生命周期管理

use crate::{Result, Error};
use crate::types::{SandboxId, SandboxSpec};
use std::process::{Command, Stdio, Child};
use std::path::PathBuf;
use serde_json::{json, Value};
use tokio::time::Duration;
use std::time::Instant;

/// Firecracker虚拟机配置
#[derive(Debug, Clone)]
pub struct VmConfig {
    /// CPU核心数
    pub vcpu_count: u32,
    /// 内存大小（MB）
    pub memory_mb: u64,
    /// 内核镜像路径
    pub kernel_image_path: PathBuf,
    /// 根文件系统路径
    pub rootfs_path: PathBuf,
    /// 引导参数
    pub boot_args: String,
}

impl From<&SandboxSpec> for VmConfig {
    fn from(spec: &SandboxSpec) -> Self {
        Self {
            vcpu_count: spec.cpu_cores,
            memory_mb: spec.memory_mb,
            kernel_image_path: PathBuf::from("/opt/koda/vmlinux.bin"),
            rootfs_path: PathBuf::from("/opt/koda/rootfs.ext4"),
            boot_args: "console=ttyS0 reboot=k panic=1 pci=off".to_string(),
        }
    }
}

/// Firecracker虚拟机
pub struct FirecrackerVm {
    /// VM ID
    pub id: SandboxId,
    /// VM规格
    pub spec: SandboxSpec,
    /// VM配置
    config: VmConfig,
    /// VM进程
    process: Option<Child>,
    /// 套接字路径
    socket_path: PathBuf,
    /// 工作目录
    work_dir: PathBuf,
}

impl FirecrackerVm {
    /// 创建新的Firecracker VM
    pub fn new(id: SandboxId, spec: SandboxSpec) -> Self {
        let socket_path = PathBuf::from(format!("/tmp/firecracker-{}.sock", id));
        let work_dir = PathBuf::from(format!("/tmp/koda-vm-{}", id));
        let config = VmConfig::from(&spec);
        
        Self {
            id,
            spec,
            config,
            process: None,
            socket_path,
            work_dir,
        }
    }
    
    /// 启动VM
    pub async fn start(&mut self) -> Result<()> {
        tracing::info!("Starting Firecracker VM: {}", self.id);
        
        // 创建工作目录
        self.create_work_directory().await?;
        
        // 生成VM配置文件
        let config_path = self.generate_vm_config().await?;
        
        // 启动Firecracker进程
        self.start_firecracker_process(&config_path).await?;
        
        // 等待VM启动完成
        self.wait_for_vm_ready().await?;
        
        tracing::info!("Firecracker VM started successfully: {}", self.id);
        Ok(())
    }
    
    /// 停止VM
    pub async fn stop(&mut self) -> Result<()> {
        if let Some(mut process) = self.process.take() {
            tracing::info!("Stopping Firecracker VM: {} (PID: {})", self.id, process.id());
            
            // 发送SIGTERM信号优雅关闭
            match process.kill() {
                Ok(()) => {
                    // 等待进程结束，由于没有异步wait，我们使用轮询
                    let start = Instant::now();
                    while start.elapsed() < Duration::from_secs(10) {
                        match process.try_wait() {
                            Ok(Some(status)) => {
                                tracing::info!("VM {} stopped with status: {}", self.id, status);
                                break;
                            }
                            Ok(None) => {
                                tokio::time::sleep(Duration::from_millis(100)).await;
                            }
                            Err(e) => {
                                tracing::error!("Error waiting for VM {} to stop: {}", self.id, e);
                                break;
                            }
                        }
                    }
                    
                    // 如果仍在运行，强制终止
                    if let Ok(None) = process.try_wait() {
                        tracing::warn!("VM {} stop timeout, forcing termination", self.id);
                        let _ = process.kill();
                    }
                }
                Err(e) => {
                    tracing::error!("Failed to kill VM {} process: {}", self.id, e);
                }
            }
        }
        
        // 清理工作目录和资源
        self.cleanup().await?;
        
        tracing::info!("Firecracker VM stopped: {}", self.id);
        Ok(())
    }
    
    /// 暂停VM
    pub async fn pause(&self) -> Result<()> {
        tracing::info!("Pausing Firecracker VM: {}", self.id);
        
        if !self.is_running() {
            return Err(Error::firecracker("VM is not running"));
        }
        
        // 通过API暂停VM
        self.send_api_request("PUT", "vm", json!({
            "state": "Paused"
        })).await?;
        
        Ok(())
    }
    
    /// 恢复VM
    pub async fn resume(&self) -> Result<()> {
        tracing::info!("Resuming Firecracker VM: {}", self.id);
        
        if !self.is_running() {
            return Err(Error::firecracker("VM is not running"));
        }
        
        // 通过API恢复VM
        self.send_api_request("PUT", "vm", json!({
            "state": "Resumed"
        })).await?;
        
        Ok(())
    }
    
    /// 检查VM是否运行
    pub fn is_running(&self) -> bool {
        // 由于借用检查问题，我们需要一个不可变的检查方法
        self.process.is_some()
    }
    
    /// 获取VM进程ID
    pub fn process_id(&self) -> Option<u32> {
        self.process.as_ref().map(|p| p.id())
    }
    
    /// 创建快照
    pub async fn create_snapshot(&self, snapshot_path: &str) -> Result<()> {
        tracing::info!("Creating snapshot for VM: {} -> {}", self.id, snapshot_path);
        
        if !self.is_running() {
            return Err(Error::firecracker("VM must be running to create snapshot"));
        }
        
        // 暂停VM以确保一致性
        self.pause().await?;
        
        // 创建快照
        let snapshot_request = json!({
            "snapshot_type": "Full",
            "snapshot_path": snapshot_path,
            "mem_file_path": format!("{}.mem", snapshot_path)
        });
        
        self.send_api_request("PUT", "snapshot/create", snapshot_request).await?;
        
        // 恢复VM
        self.resume().await?;
        
        tracing::info!("Snapshot created successfully for VM: {}", self.id);
        Ok(())
    }
    
    /// 从快照恢复
    pub async fn restore_from_snapshot(&mut self, snapshot_path: &str) -> Result<()> {
        tracing::info!("Restoring VM: {} from {}", self.id, snapshot_path);
        
        // 验证快照文件存在
        if !PathBuf::from(snapshot_path).exists() {
            return Err(Error::firecracker(format!("Snapshot file not found: {}", snapshot_path)));
        }
        
        // 如果VM正在运行，先停止
        if self.is_running() {
            self.stop().await?;
        }
        
        // 创建工作目录
        self.create_work_directory().await?;
        
        // 生成恢复配置
        let config_path = self.generate_restore_config(snapshot_path).await?;
        
        // 启动Firecracker进程并恢复快照
        self.start_firecracker_process(&config_path).await?;
        
        // 等待VM恢复完成
        self.wait_for_vm_ready().await?;
        
        tracing::info!("VM restored successfully from snapshot: {}", snapshot_path);
        Ok(())
    }
    
    /// 创建工作目录
    async fn create_work_directory(&self) -> Result<()> {
        tokio::fs::create_dir_all(&self.work_dir).await
            .map_err(|e| Error::firecracker(format!("Failed to create work directory: {}", e)))?;
        Ok(())
    }
    
    /// 生成VM配置文件
    async fn generate_vm_config(&self) -> Result<PathBuf> {
        let config_path = self.work_dir.join("vm_config.json");
        
        let config_json = json!({
            "boot-source": {
                "kernel_image_path": self.config.kernel_image_path,
                "boot_args": self.config.boot_args
            },
            "drives": [{
                "drive_id": "rootfs",
                "path_on_host": self.config.rootfs_path,
                "is_root_device": true,
                "is_read_only": false
            }],
            "machine-config": {
                "vcpu_count": self.config.vcpu_count,
                "mem_size_mib": self.config.memory_mb
            },
            "logger": {
                "log_path": self.work_dir.join("firecracker.log"),
                "level": "Info"
            }
        });
        
        tokio::fs::write(&config_path, serde_json::to_string_pretty(&config_json)?)
            .await
            .map_err(|e| Error::firecracker(format!("Failed to write config file: {}", e)))?;
        
        Ok(config_path)
    }
    
    /// 生成恢复配置文件
    async fn generate_restore_config(&self, snapshot_path: &str) -> Result<PathBuf> {
        let config_path = self.work_dir.join("restore_config.json");
        
        let config_json = json!({
            "snapshot-path": snapshot_path,
            "mem-file-path": format!("{}.mem", snapshot_path),
            "enable-diff-snapshots": false,
            "resume-vm": true
        });
        
        tokio::fs::write(&config_path, serde_json::to_string_pretty(&config_json)?)
            .await
            .map_err(|e| Error::firecracker(format!("Failed to write restore config: {}", e)))?;
        
        Ok(config_path)
    }
    
    /// 启动Firecracker进程
    async fn start_firecracker_process(&mut self, config_path: &PathBuf) -> Result<()> {
        let mut cmd = Command::new("firecracker");
        cmd.arg("--api-sock").arg(&self.socket_path)
           .arg("--config-file").arg(config_path)
           .stdout(Stdio::piped())
           .stderr(Stdio::piped());
        
        let process = cmd.spawn()
            .map_err(|e| Error::firecracker(format!("Failed to start Firecracker: {}", e)))?;
        
        self.process = Some(process);
        
        // 等待套接字文件创建
        self.wait_for_socket().await?;
        
        Ok(())
    }
    
    /// 等待套接字文件创建
    async fn wait_for_socket(&self) -> Result<()> {
        for _ in 0..50 { // 最多等待5秒
            if self.socket_path.exists() {
                return Ok(());
            }
            tokio::time::sleep(Duration::from_millis(100)).await;
        }
        
        Err(Error::firecracker("Timeout waiting for Firecracker socket"))
    }
    
    /// 等待VM启动完成
    async fn wait_for_vm_ready(&self) -> Result<()> {
        // 发送启动请求
        self.send_api_request("PUT", "actions", json!({
            "action_type": "InstanceStart"
        })).await?;
        
        // 等待VM进入运行状态
        for _ in 0..100 { // 最多等待10秒
            match self.get_vm_state().await {
                Ok(state) if state == "Running" => {
                    return Ok(());
                }
                Ok(_) => {
                    tokio::time::sleep(Duration::from_millis(100)).await;
                }
                Err(_) => {
                    tokio::time::sleep(Duration::from_millis(100)).await;
                }
            }
        }
        
        Err(Error::firecracker("Timeout waiting for VM to become ready"))
    }
    
    /// 获取VM状态
    async fn get_vm_state(&self) -> Result<String> {
        let response = self.send_api_request("GET", "vm", json!({})).await?;
        Ok(response.get("state")
            .and_then(|s| s.as_str())
            .unwrap_or("Unknown")
            .to_string())
    }
    
    /// 发送API请求到Firecracker
    async fn send_api_request(&self, method: &str, path: &str, body: Value) -> Result<Value> {
        // 这里应该实现真正的HTTP over Unix socket请求
        // 为了POC，我们模拟一个成功的响应
        tracing::debug!("Sending {} request to {}: {}", method, path, body);
        
        // 模拟API调用延迟
        tokio::time::sleep(Duration::from_millis(10)).await;
        
        // 根据请求路径返回模拟响应
        match path {
            "vm" if method == "GET" => Ok(json!({"state": "Running"})),
            "vm" if method == "PUT" => Ok(json!({"status": "success"})),
            "actions" => Ok(json!({"status": "success"})),
            "snapshot/create" => Ok(json!({"status": "success"})),
            _ => Ok(json!({"status": "success"}))
        }
    }
    
    /// 清理资源
    async fn cleanup(&self) -> Result<()> {
        // 删除套接字文件
        if self.socket_path.exists() {
            if let Err(e) = tokio::fs::remove_file(&self.socket_path).await {
                tracing::warn!("Failed to remove socket file: {}", e);
            }
        }
        
        // 清理工作目录
        if self.work_dir.exists() {
            if let Err(e) = tokio::fs::remove_dir_all(&self.work_dir).await {
                tracing::warn!("Failed to remove work directory: {}", e);
            }
        }
        
        Ok(())
    }
}

impl Drop for FirecrackerVm {
    fn drop(&mut self) {
        if self.is_running() {
            tracing::warn!("VM {} dropped while still running", self.id);
            // 在实际实现中，这里应该异步停止VM
            // 由于Drop trait不支持async，我们只能记录警告
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{SandboxSpec, Language};
    use uuid::Uuid;
    
    #[tokio::test]
    async fn test_vm_lifecycle() {
        let id = Uuid::new_v4();
        let spec = SandboxSpec {
            language: Language::Python,
            ..Default::default()
        };
        
        let mut vm = FirecrackerVm::new(id, spec);
        
        assert!(!vm.is_running());
        assert_eq!(vm.process_id(), None);
        
        // 注意：由于我们不能在测试中真正启动Firecracker，
        // 这些测试会因为缺少Firecracker二进制而失败
        // 在真实环境中需要有Firecracker可执行文件
    }
    
    #[test]
    fn test_vm_creation() {
        let id = Uuid::new_v4();
        let spec = SandboxSpec::default();
        
        let vm = FirecrackerVm::new(id, spec.clone());
        
        assert_eq!(vm.id, id);
        assert_eq!(vm.spec, spec);
        assert!(!vm.is_running());
    }
    
    #[test]
    fn test_vm_config_from_spec() {
        let spec = SandboxSpec {
            language: Language::Python,
            cpu_cores: 2,
            memory_mb: 512,
            ..Default::default()
        };
        
        let config = VmConfig::from(&spec);
        assert_eq!(config.vcpu_count, 2);
        assert_eq!(config.memory_mb, 512);
        assert!(config.boot_args.contains("console=ttyS0"));
    }
    
    #[tokio::test]
    async fn test_work_directory_creation() {
        let id = Uuid::new_v4();
        let spec = SandboxSpec::default();
        let vm = FirecrackerVm::new(id, spec);
        
        // 测试工作目录创建
        let result = vm.create_work_directory().await;
        assert!(result.is_ok());
        
        // 清理测试目录
        if vm.work_dir.exists() {
            let _ = tokio::fs::remove_dir_all(&vm.work_dir).await;
        }
    }
}